{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KSSNwAL63DBc"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g5Ju84HQ3DBe"
   },
   "source": [
    "If the `econml` and `wget` python packages are not installed on your machine install them by running the cells below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install econml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q18wYNnk3DBg"
   },
   "outputs": [],
   "source": [
    "!pip install wget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m6IN5aVT3DBg"
   },
   "source": [
    "Importing all the necessary components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "His1WYC83DBg"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "import scipy.special\n",
    "from sklearn.linear_model import LassoCV, LinearRegression\n",
    "from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor\n",
    "from sklearn.tree import plot_tree\n",
    "from statsmodels.api import OLS\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.model_selection import cross_val_predict\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Uud5pCem3DBh"
   },
   "outputs": [],
   "source": [
    "import wget\n",
    "import os\n",
    "\n",
    "if os.path.exists('datasets.py'):\n",
    "    os.remove('datasets.py')\n",
    "wget.download('https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/T/datasets.py')\n",
    "\n",
    "if os.path.exists('myxgb.py'):\n",
    "    os.remove('myxgb.py')\n",
    "wget.download('https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/T/myxgb.py')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hdo6-SJT3DBh"
   },
   "source": [
    "# Setting the High Level Parameters for the Notebook\n",
    "\n",
    "- The 401k dataset is downloaded from the source by the code and no need to further download anything:\n",
    "https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/data/401k.csv\n",
    "\n",
    "- The welfare dataset is downloaded from the source by the code and no need to further download anything:\n",
    "https://github.com/gsbDBI/ExperimentData/blob/master/Welfare/ProcessedData/welfarenolabel3.csv.\n",
    "It is drawn from the analysis in this paper: [Green and Kern, 2012, Modeling Heterogeneous Treatment Effects in Survey Experiments with Bayesian Additive Regression Trees](https://github.com/gsbDBI/ExperimentData/blob/master/Welfare/Green%20and%20Kern%20BART.pdf)\n",
    "\n",
    "- The Criteo dataset has to be downloaded from here:\n",
    "https://www.kaggle.com/code/hughhuyton/criteo-uplift-modelling/input. This file should be downloaded and stored in the same folder as the notebook with the name `criteo-uplift-v2.1.csv`. Warning that this dataset is quite large and the notebook will take longer to execute than the other datasets. If you are using Google Colab, the best way to bring this dataset to your local folder is the following:\n",
    "    - Go to the Kaggle url above and download the dataset\n",
    "    - Upload the zipped dataset to a folder in your google drive e.g. `path-to-zip/criteo-uplift-v2.1.csv.zip`\n",
    "    - Then run the following commands on colab, which will mount your google drive and will unzip the dataset from the drive to the current folder on the local colab machine\n",
    "    ```python\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/drive')\n",
    "    \n",
    "    !unzip -q '/content/drive/My Drive/path-to-zip/criteo-uplift-v2.1.csv.zip' -d .\n",
    "    ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JQGC_iwe3DBi"
   },
   "outputs": [],
   "source": [
    "dataset = 'welfare'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mw-jBrpA3DBi"
   },
   "outputs": [],
   "source": [
    "if dataset == '401k':\n",
    "    verbose = 0  # verbosity of auto-ml\n",
    "    n_splits = 5  # cross-fitting and cross-validation splits\n",
    "    cfit = False\n",
    "    data = '401k'  # which dataset, one of {'401k', 'criteo', 'welfare', 'poverty', 'charity'}\n",
    "    plot = True  # whether to plot results\n",
    "    # feature to use as x axis in plotting, e.g. for criteo 'f1', for 401k 'inc', for welfare 'polviews'\n",
    "    xfeat = 'inc'\n",
    "    # Formula for the BLP of CATE regression.\n",
    "    blp_formula = 'np.log(inc)'  # e.g. 'f1' for criteo, np.log(inc)' for 401k, 'C(polviews)' for the welfare case.\n",
    "    blp_formula_short = 'log(inc)'\n",
    "    blp_formula2 = 'np.log(inc) + np.power(np.log(inc), 2) + np.power(np.log(inc), 3) + np.power(np.log(inc), 4)'\n",
    "    blp_formula2_short = 'poly(log(inc), 4)'\n",
    "    control_feats = 'all'  # list of subset of features to be used as confounders or the string 'all' for everything\n",
    "    # list of subset of control features to be used for CATE model or the string 'all' for all controls\n",
    "    hetero_feats = 'all'\n",
    "    cov_clip = .01  # clipping of treatment variance p(x)*(1-p(x)), whenever used in inverse propensities\n",
    "    binary_y = False\n",
    "    random_seed = 1\n",
    "\n",
    "    # cost of treatment when performing optimal policy learning, can also be viewed as \"threshold for treatment\"\n",
    "    treatment_cost = 4000\n",
    "\n",
    "    # For semi-synthetic data generation\n",
    "    semi_synth = False  # Whether true outcome y should be replaced by a fake outcome from a known CEF\n",
    "    simple_synth = True  # Whether the true CEF of the fake y should be simple or fitted from data\n",
    "    max_depth = 2  # max depth of random forest during for semi-synthetic model fitting\n",
    "    scale = .2  # magnitude of noise in semi-synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ygfUjoyF3DBi"
   },
   "outputs": [],
   "source": [
    "if dataset == 'criteo':\n",
    "    verbose = 0  # verbosity of auto-ml\n",
    "    # cross-fitting and cross-validation splits; increase to 5 for better statistical performance\n",
    "    # albeit notebook will take much longer to execute\n",
    "    n_splits = 2\n",
    "    cfit = False\n",
    "    data = 'criteo'  # which dataset, one of {'401k', 'criteo', 'welfare', 'poverty', 'charity'}\n",
    "    plot = True  # whether to plot results\n",
    "    # feature to use as x axis in plotting, e.g. for criteo 'f1', for 401k 'inc', for welfare 'polviews'\n",
    "    xfeat = 'f3'\n",
    "    control_feats = 'all'  # list of subset of features to be used as confounders or the string 'all' for everything\n",
    "    # list of subset of control features to be used for CATE model or the string 'all' for all controls\n",
    "    hetero_feats = 'all'\n",
    "    cov_clip = .01  # clipping of treatment variance p(x)*(1-p(x)), whenever used in inverse propensities\n",
    "    binary_y = True\n",
    "    random_seed = 1\n",
    "    # Formula for the BLP of CATE regression.\n",
    "    blp_formula = 'f3'  # e.g. 'f1' for criteo, np.log(inc)' for 401k, 'C(polviews)' for the welfare case.\n",
    "    blp_formula_short = 'f3'\n",
    "    blp_formula2 = 'np.power(f3, 2)'\n",
    "    blp_formula2_short = 'poly(f3, 2)'\n",
    "    # cost of treatment when performing optimal policy learning, can also be viewed as \"threshold for treatment\"\n",
    "    treatment_cost = 0.04\n",
    "\n",
    "    # For semi-synthetic data generation\n",
    "    semi_synth = False  # Whether true outcome y should be replaced by a fake outcome from a known CEF\n",
    "    simple_synth = True  # Whether the true CEF of the fake y should be simple or fitted from data\n",
    "    max_depth = 2  # max depth of random forest during for semi-synthetic model fitting\n",
    "    scale = .2  # magnitude of noise in semi-synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "v-OnXkur3DBi"
   },
   "outputs": [],
   "source": [
    "if dataset == 'welfare':\n",
    "    verbose = 0  # verbosity of auto-ml\n",
    "    n_splits = 5  # cross-fitting and cross-validation splits\n",
    "    cfit = True\n",
    "    data = 'welfare'  # which dataset, one of {'401k', 'criteo', 'welfare', 'poverty', 'star'}\n",
    "    plot = True  # whether to plot results\n",
    "    # feature to use as x axis in plotting, e.g. for criteo 'f1', for 401k 'inc', for welfare 'polviews'\n",
    "    xfeat = 'polviews'\n",
    "    # Formula for the BLP of CATE regression.\n",
    "    blp_formula = 'C(polviews)'  # e.g. 'f1' for criteo, np.log(inc)' for 401k, 'C(polviews)' for the welfare case.\n",
    "    blp_formula_short = 'C(polviews)'\n",
    "    blp_formula2 = 'polviews'\n",
    "    blp_formula2_short = 'polviews'\n",
    "    control_feats = 'all'  # list of subset of features to be used as confounders or the string 'all' for everything\n",
    "    # list of subset of control features to be used for CATE model or the string 'all' for all controls\n",
    "    hetero_feats = 'all'\n",
    "    cov_clip = .01  # clipping of treatment variance p(x)*(1-p(x)), whenever used in inverse propensities\n",
    "    binary_y = True\n",
    "    random_seed = 1\n",
    "\n",
    "    # cost of treatment when performing optimal policy learning, can also be viewed as \"threshold for treatment\"\n",
    "    treatment_cost = -.3\n",
    "\n",
    "    # For semi-synthetic data generation\n",
    "    semi_synth = False  # Whether true outcome y should be replaced by a fake outcome from a known CEF\n",
    "    simple_synth = True  # Whether the true CEF of the fake y should be simple or fitted from data\n",
    "    max_depth = 2  # max depth of random forest during for semi-synthetic model fitting\n",
    "    scale = .2  # magnitude of noise in semi-synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simple_true_cef(D, X):  # simple CEF of the outcome for semi-synthetic data\n",
    "    return .5 * np.array(X)[:, 1] * D + np.array(X)[:, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ru6z3CV93DBi"
   },
   "source": [
    "# Constructing the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zeJf9EfT3DBj"
   },
   "outputs": [],
   "source": [
    "from datasets import fetch_data_generator\n",
    "\n",
    "get_data, abtest, true_cef, true_cate = fetch_data_generator(data=data, semi_synth=semi_synth,\n",
    "                                                             simple_synth=simple_synth,\n",
    "                                                             scale=scale, true_f=simple_true_cef,\n",
    "                                                             max_depth=max_depth)\n",
    "X, D, y, groups = get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bcCu7Hc03DBj"
   },
   "outputs": [],
   "source": [
    "if semi_synth:\n",
    "    true_ate = np.mean(true_cate(X))\n",
    "    print(f'True ATE: {true_ate}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MFCbmYAy3DBj"
   },
   "outputs": [],
   "source": [
    "def rmse(cate, preds):\n",
    "    return np.sqrt(np.mean((cate - preds)**2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dtSQ6hcm3DBj"
   },
   "source": [
    "# Data Analysis\n",
    "\n",
    "We now have our data $X$, $D$, $y$, of controls, treatments and outcomes. In some datasets, we also have \"groups\", also known as \"clusters\". These are group ids, that define a group of samples that are believed to be correlated through unobesrved factors. For instance, in randomized experiments when a whole class is being treated and we have data at the student level, the students in a class constitute a cluster, as their outcome variables are most probably correlated. In such settings, it is helpful to account for the cluster correlations when calculating confidence intervals and when performing sample splitting for either cross-validation or for nuisance estimation. This notebook will not deal with such group clusters. A more tailored analysis is required.\n",
    "\n",
    "We will be assuming throughout that conditional ignorability is satisfied if we control for all the variables $X$, i.e. the potential outcomes $Y(1), Y(0)$ satisfy\n",
    "\\begin{align}\n",
    "Y(1), Y(0) ~\\perp\\hspace{-1em}\\perp~D \\mid X\n",
    "\\end{align}\n",
    "Equivalently, we assume that the DAG the corresponds to our setting satisfies that $X$ is a valid adjustment set between $D$ and $Y$, i.e. it blocks all backdoor paths in the DAG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MdcaOSBN3DBj"
   },
   "outputs": [],
   "source": [
    "X.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PSG5kH07egGB"
   },
   "outputs": [],
   "source": [
    "if control_feats != 'all':\n",
    "    X = X[control_feats]\n",
    "    print(X.describe())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FQC_L8BamvSz"
   },
   "outputs": [],
   "source": [
    "X = X - X.mean(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5PYonx-43DBj"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 5))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.hist(D)\n",
    "plt.title('Treatment')\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.hist(y)\n",
    "plt.title('Outcome')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8l30d3fF3DBj"
   },
   "outputs": [],
   "source": [
    "# simple two means estimate which would be wrong unless an randomized trial\n",
    "OLS(y, np.hstack([np.ones((D.shape[0], 1)), D.reshape(-1, 1)])).fit(cov_type='HC1').summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5KWdr4723DBj"
   },
   "source": [
    "# Split Train and Validation and Test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cr-geLJT3DBj"
   },
   "source": [
    "The training data will be used to fit the various CATE models. The validation data will be used for scoring and selection of the best CATE model or best ensemble of CATE models. The test data will be used for testing and evaluation of the performance of the best chosen model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QHSG1Abx3DBj"
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import StratifiedShuffleSplit\n",
    "\n",
    "DX = np.column_stack((D, X))\n",
    "labels = np.column_stack((D, y)) if binary_y else D\n",
    "train, val = next(StratifiedShuffleSplit(n_splits=2, train_size=.6,\n",
    "                                         random_state=random_seed).split(DX, labels))\n",
    "X, Xval, D, Dval, y, yval = X.iloc[train], X.iloc[val], D[train], D[val], y[train], y[val]\n",
    "\n",
    "DXval = np.column_stack((Dval, Xval))\n",
    "labels_val = np.column_stack((Dval, yval)) if binary_y else Dval\n",
    "val, test = next(StratifiedShuffleSplit(n_splits=2, train_size=.6,\n",
    "                                        random_state=random_seed).split(DXval, labels_val))\n",
    "Xval, Xtest, Dval, Dtest, yval, ytest = Xval.iloc[val], Xval.iloc[test], Dval[val], Dval[test], yval[val], yval[test]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wQJBSfsP3DBj"
   },
   "source": [
    "# Nuisance Cross-Fitted Estimation and Prediction\n",
    "\n",
    "We will estimate regression models for each of the nuisance functions that arise in CATE learning approaches. The five models correspond to the following five predictive problems:\n",
    "\\begin{align}\n",
    "\\text{model_y} ~\\rightarrow~& q(x) := E[Y\\mid X=x]\\\\\n",
    "\\text{model_t} ~\\rightarrow~& p(x) := E[D\\mid X=x] = \\Pr(D=1\\mid X=x)\\\\\n",
    "\\text{model_reg_zero} ~\\rightarrow~& g_0(x) := E[Y\\mid D=0, X=x]\\\\\n",
    "\\text{model_reg_one} ~\\rightarrow~& g_1(x) := E[Y\\mid D=1, X=x]\\\\\n",
    "\\end{align}\n",
    "We will use gradient boosting regression with early stopping for each of these models.\n",
    "\n",
    "For each of the nuisance models we perform cross-fitting to get out-of-fold predictions from each of these nuisance models. At the end of this process, we will have for each sample $i$, the following out-of-fold nuisance values:\n",
    "\\begin{align}\n",
    "\\text{reg_preds_t} \\rightarrow~& \\hat{g}_0(X_i) (1 - D_i) + \\hat{g}_1(X_i) D_i &\n",
    "\\text{reg_one_preds_t} \\rightarrow~& \\hat{g}_1(X_i) &\n",
    "\\text{reg_zero_preds_t} \\rightarrow~& \\hat{g}_0(X_i)\\\\\n",
    "\\text{res_preds} \\rightarrow~& \\hat{q}(X_i) &\n",
    "\\text{prop_preds} \\rightarrow~& \\hat{p}(X_i)\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IDYPiTP43DBk"
   },
   "outputs": [],
   "source": [
    "from myxgb import xgb_reg, xgb_clf, xgb_wreg, RegWrapper\n",
    "\n",
    "\n",
    "def auto_reg():\n",
    "    return xgb_reg(random_seed)\n",
    "\n",
    "\n",
    "# Disclaimer: The remainder of the code assumes that the `auto_clf` model returns\n",
    "# the probability of class 1, when one calls auto_clf().predict(X)\n",
    "# and not the 0/1 classification. This is what the xgb_clf class does.\n",
    "def auto_clf():\n",
    "    return RegWrapper(xgb_clf(random_seed))\n",
    "\n",
    "\n",
    "def auto_weighted_reg():\n",
    "    return xgb_wreg(random_seed)\n",
    "\n",
    "\n",
    "modely = auto_clf if binary_y else auto_reg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8zlXri4Z3DBk"
   },
   "outputs": [],
   "source": [
    "if cfit:\n",
    "    cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_seed)\n",
    "    stratification_label = 2 * D + y if binary_y else D\n",
    "    splits = list(cv.split(X, stratification_label))\n",
    "else:\n",
    "    splits = [(np.arange(X.shape[0]), np.arange(X.shape[0]))]\n",
    "\n",
    "n = X.shape[0]\n",
    "reg_preds = np.zeros(n)\n",
    "reg_zero_preds = np.zeros(n)\n",
    "reg_one_preds = np.zeros(n)\n",
    "reg_preds_t = np.zeros(n)\n",
    "reg_zero_preds_t = np.zeros(n)\n",
    "reg_one_preds_t = np.zeros(n)\n",
    "\n",
    "DX = np.column_stack((D, X))\n",
    "for train, test in splits:\n",
    "    reg = modely().fit(DX[train], y[train])\n",
    "    reg_preds[test] = reg.predict(DX[test])\n",
    "    reg_one_preds[test] = reg.predict(np.column_stack([np.ones(len(test)), X.iloc[test]]))\n",
    "    reg_zero_preds[test] = reg.predict(np.column_stack([np.zeros(len(test)), X.iloc[test]]))\n",
    "\n",
    "    reg_zero = modely().fit(X.iloc[train][D[train] == 0], y[train][D[train] == 0])\n",
    "    reg_one = modely().fit(X.iloc[train][D[train] == 1], y[train][D[train] == 1])\n",
    "    reg_zero_preds_t[test] = reg_zero.predict(X.iloc[test])\n",
    "    reg_one_preds_t[test] = reg_one.predict(X.iloc[test])\n",
    "    reg_preds_t[test] = reg_zero_preds_t[test] * (1 - D[test]) + reg_one_preds_t[test] * D[test]\n",
    "\n",
    "res_preds = cross_val_predict(modely(), X, y, cv=splits)\n",
    "prop_preds = cross_val_predict(auto_clf(), X, D, cv=splits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0UOAQK8-3DBk"
   },
   "source": [
    "# Evaluating Nuisance Model Performance\n",
    "\n",
    "We now also evaluate the performance of the selected models in terms of R^2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tHjDHFEc3DBk"
   },
   "outputs": [],
   "source": [
    "def r2score(y, ypred):\n",
    "    return 1 - np.mean((y - ypred)**2) / np.var(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ht_oIAkc3DBk"
   },
   "outputs": [],
   "source": [
    "print(f\"R^2 of model for (y ~ X): {r2score(y, res_preds):.4f}\")\n",
    "print(f\"R^2 of model for (D ~ X): {r2score(D, prop_preds):.4f}\")\n",
    "print(f\"R^2 of two-learner model for (y ~ X | D==0): {r2score(y[D == 0], reg_zero_preds_t[D == 0]):.4f}\")\n",
    "print(f\"R^2 of two-learner model for (y ~ X | D==1): {r2score(y[D == 1], reg_one_preds_t[D == 1]):.4f}\")\n",
    "print(f\"R^2 of two-learner model for (y ~ D, X): {r2score(y, reg_preds_t):.4f}\")\n",
    "print(f\"R^2 of single-learner model for (y ~ D, X): {r2score(y, reg_preds):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mn78GUwL3DBk"
   },
   "source": [
    "# Doubly-Robust ATE Estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KeE7awhz3DBk"
   },
   "source": [
    "Using the doubly robust method. In particular, we construct the doubly robust variables\n",
    "\\begin{align}\n",
    "Y_i^{DR}(\\hat{g},\\hat{p}) := \\hat{g}_1(X_i) - \\hat{g}_0(X_i) + (Y_i - \\hat{g}_{D_i}(X_i))\\frac{D_i - \\hat{p}(X_i)}{\\hat{p}(X_i) (1-\\hat{p}(X_i))}\n",
    "\\end{align}\n",
    "and then we estimate:\n",
    "\\begin{align}\n",
    "ATE = E_n\\left[Y^{DR}(\\hat{g},\\hat{p})\\right]\n",
    "\\end{align}\n",
    "This should be more efficient in the worst-case and should be returning a consistent estimate of the ATE even beyond RCTs and will also correctly account for any imbalances or violations of the randomization assumption in an RCT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b1QMzWtv3DBk"
   },
   "outputs": [],
   "source": [
    "dr_preds = reg_one_preds_t - reg_zero_preds_t\n",
    "dr_preds += (y - reg_preds_t) * (D - prop_preds) / np.clip(prop_preds * (1 - prop_preds), cov_clip, np.inf)\n",
    "\n",
    "OLS(dr_preds, np.ones((len(dr_preds), 1))).fit(cov_type='HC1').summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ShzvhOGY3DBk"
   },
   "outputs": [],
   "source": [
    "from statsmodels.formula.api import ols\n",
    "df = X.copy()\n",
    "df['dr'] = dr_preds\n",
    "lr = ols('dr ~ ' + blp_formula, df).fit(cov_type='HC1')\n",
    "lr.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A0cSzR2x3DBk"
   },
   "outputs": [],
   "source": [
    "grid = np.unique(np.percentile(X[xfeat], np.arange(0, 102, 2)))\n",
    "Xpd = pd.DataFrame(np.tile(np.median(X, axis=0, keepdims=True), (len(grid), 1)),\n",
    "                   columns=X.columns)\n",
    "Xpd[xfeat] = grid\n",
    "pred_df = lr.get_prediction(Xpd).summary_frame(alpha=.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X1ep7hRS3DBk"
   },
   "outputs": [],
   "source": [
    "plt.plot(Xpd[xfeat], pred_df['mean'])\n",
    "plt.fill_between(Xpd[xfeat], pred_df['mean_ci_lower'], pred_df['mean_ci_upper'], alpha=.4)\n",
    "plt.xlabel(xfeat + ' (other features fixed at median value)')\n",
    "plt.title('Predicted CATE BLP: cate ~' + blp_formula)\n",
    "plt.ylabel('CATE')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P-9GcXSk3DBk"
   },
   "outputs": [],
   "source": [
    "from statsmodels.formula.api import ols\n",
    "df = X.copy()\n",
    "df['dr'] = dr_preds\n",
    "lr = ols('dr ~ ' + blp_formula2, df).fit(cov_type='HC1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9OFr9AqC3DBl"
   },
   "outputs": [],
   "source": [
    "grid = np.unique(np.percentile(X[xfeat], np.arange(0, 102, 2)))\n",
    "Xpd = pd.DataFrame(np.tile(np.median(X, axis=0, keepdims=True), (len(grid), 1)),\n",
    "                   columns=X.columns)\n",
    "Xpd[xfeat] = grid\n",
    "pred_df2 = lr.get_prediction(Xpd).summary_frame(alpha=.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_nncoMd03DBl"
   },
   "outputs": [],
   "source": [
    "plt.plot(Xpd[xfeat], pred_df2['mean'])\n",
    "plt.fill_between(Xpd[xfeat], pred_df2['mean_ci_lower'], pred_df2['mean_ci_upper'], alpha=.4)\n",
    "plt.xlabel(xfeat + ' (other features fixed at median value)')\n",
    "plt.ylabel('CATE')\n",
    "plt.title('cate ~' + blp_formula2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wwdl2CiJ3DBl"
   },
   "source": [
    "# CATE Model Estimation with Meta-Learners"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "puSsJBmh3DBp"
   },
   "source": [
    "We specify which indices of the X variables we want to use for heterogeneity. Let's denote these subset of variables with $Z$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BCLIIZ7W3DBp"
   },
   "outputs": [],
   "source": [
    "if hetero_feats == 'all':\n",
    "    hetero_feats = X.columns\n",
    "Z, Zval, Ztest = X[hetero_feats], Xval[hetero_feats], Xtest[hetero_feats]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yzjQhSxB3DBq"
   },
   "source": [
    "We specify a generic automl approach for training the final CATE model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q5Q0y9WN3DBq"
   },
   "source": [
    "### Single Learner (S-Learner)\n",
    "\n",
    "The S-Learner simply trains a model to predict:\n",
    "\\begin{align}\n",
    "\\hat{g}(1, X_i) - \\hat{g}(0, X_i) \\sim Z_i\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rK9X06wY3DBq"
   },
   "outputs": [],
   "source": [
    "# slearner\n",
    "slearner = auto_reg().fit(Z, reg_one_preds - reg_zero_preds)\n",
    "slearner_cates = slearner.predict(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_CR_Utvh3DBq"
   },
   "source": [
    "### Two Learner (T-Learner)\n",
    "\n",
    "The T-Learner simply trains a model to predict:\n",
    "\\begin{align}\n",
    "\\hat{g}_1(X_i) - \\hat{g}_0(X_i) \\sim Z_i\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IbMnTqwY3DBq"
   },
   "outputs": [],
   "source": [
    "# tlearner\n",
    "tlearner = auto_reg().fit(Z, reg_one_preds_t - reg_zero_preds_t)\n",
    "tlearner_cates = tlearner.predict(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k3e2-_em3DBq"
   },
   "source": [
    "### Cross Learner (X-Learner)\n",
    "\n",
    "The X-Learner simply trains a two treatment effect models using all the variables $X$:\n",
    "\\begin{align}\n",
    "\\hat{\\tau}_1: Y_i - \\hat{g}_0(X_i) \\sim X_i \\\\\n",
    "\\hat{\\tau}_0: \\hat{g}_1(X_i) - Y_i \\sim X_i\n",
    "\\end{align}\n",
    "And then for each target sample $X_i$ we select a mixture based on the propensity.\n",
    "If $D=1$ is more probable, then we use more heavily $\\hat{\\tau}_0$, since that uses the model $\\hat{g}_1$ which was trained on more data similar to $X_i$ (as indicated by the propensity). Similarly, if $D=0$ is more probably we put more weight on $\\hat{\\tau}_1$.\n",
    "\\begin{align}\n",
    "\\hat{\\tau}(X_i) = \\hat{p}(X_i) \\hat{\\tau}_0(X_i) + (1 - \\hat{p}(X_i)) \\hat{\\tau}_1(X_i)\n",
    "\\tag{xtarget}\n",
    "\\end{align}\n",
    "Finally, to learn a CATE model that only depends on $Z$, we predict the CATEs $\\hat{\\tau}$ from $Z$:\n",
    "\\begin{align}\n",
    "\\hat{\\tau}(X_i) \\sim Z_i\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "srD1loPF3DBq"
   },
   "outputs": [],
   "source": [
    "# xlearner\n",
    "tau1_preds = y[D == 1] - reg_zero_preds_t[D == 1]\n",
    "tau0_preds = reg_one_preds_t[D == 0] - y[D == 0]\n",
    "tau1 = auto_reg().fit(X[D == 1], tau1_preds)\n",
    "tau0 = auto_reg().fit(X[D == 0], tau0_preds)\n",
    "xtarget = prop_preds * tau0.predict(X) + (1 - prop_preds) * tau1.predict(X)\n",
    "xlearner = auto_reg().fit(Z, xtarget)\n",
    "xlearner_cates = xlearner.predict(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dr48wCTr3DBq"
   },
   "source": [
    "### Doubly Robust Learner (DR-Learner)\n",
    "\n",
    "The DR-Learner (in particular, the variant based on the T-Learner) adds a de-biasing correction to the T-Learner using the propensity. In particular, we construct the doubly robust variables\n",
    "\\begin{align}\n",
    "Y_i^{DR}(\\hat{g},\\hat{p}) := \\hat{g}_1(X_i) - \\hat{g}_0(X_i) + (Y_i - \\hat{g}_{D_i}(X_i))\\frac{D_i - \\hat{p}(X_i)}{\\hat{p}(X_i) (1-\\hat{p}(X_i))}\n",
    "\\end{align}\n",
    "and then we train a CATE model by predicting these variables from $Z_i$:\n",
    "\\begin{align}\n",
    "Y_i^{DR}(\\hat{g},\\hat{p}) \\sim Z_i\n",
    "\\end{align}\n",
    "For stability we clip the co-variance that appears in the denominator at some value $c$ bounded away from zero, i.e.\n",
    "\\begin{align}\n",
    "Y_i^{DR}(\\hat{g},\\hat{p}) := \\hat{g}_1(X_i) - \\hat{g}_0(X_i) + (Y_i - \\hat{g}_{D_i}(X_i))\\frac{D_i - \\hat{p}(X_i)}{\\min\\{c, \\hat{p}(X_i) (1-\\hat{p}(X_i))\\}}\n",
    "\\tag{dr_preds}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rJFgOCA_3DBq"
   },
   "outputs": [],
   "source": [
    "# drlearner\n",
    "dr_preds = reg_one_preds_t - reg_zero_preds_t\n",
    "dr_preds += (y - reg_preds_t) * (D - prop_preds) / np.clip(prop_preds * (1 - prop_preds), cov_clip, np.inf)\n",
    "drlearner = auto_reg().fit(Z, dr_preds)\n",
    "drlearner_cates = drlearner.predict(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y0vokZMw3DBq"
   },
   "source": [
    "### Residual Learner (R-Learner)\n",
    "\n",
    "The R-Learner uses the fact that if the true CATE model\n",
    "$$\\theta(X)=E[Y(1)-Y(0)\\mid X]$$\n",
    "only depends on the subset of variables $Z$, then we can write:\n",
    "\\begin{align}\n",
    "Y =~& \\tau(Z) D + b(X) + \\epsilon, & E[\\epsilon\\mid D, X]=~& 0\n",
    "\\end{align}\n",
    "This then implies that we can also write:\n",
    "\\begin{align}\n",
    "Y - E[Y\\mid X] =~& \\tau(Z) (D - E[D|X]) + \\epsilon, & E[\\epsilon\\mid D, X]=~& 0\n",
    "\\end{align}\n",
    "Hence, if we define the residual outcome and the residual treatment,\n",
    "\\begin{align}\n",
    "\\tilde{Y} :=~& Y - E[Y\\mid X] \\tag{yres}\\\\\n",
    "\\tilde{D} :=~& D - E[D\\mid X] \\tag{Dres}\n",
    "\\end{align}\n",
    "then we have:\n",
    "\\begin{align}\n",
    "\\tilde{Y} =~& \\tau(Z) \\tilde{D} + \\epsilon, & E[\\epsilon\\mid D, X]=~& 0\n",
    "\\end{align}\n",
    "Thus we can estimate the CATE by predicting $\\tilde{Y}$ from $\\tilde{D}, Z$ with a separable model of the form $\\tau(Z) \\tilde{D}$. This also implies that $\\tau$ is the minimizer of the square loss:\n",
    "\\begin{align}\n",
    "E\\left[ \\left(\\tilde{Y} - \\tau(Z) \\tilde{D}\\right)^2 \\right]\n",
    "\\end{align}\n",
    "This loss can also be re-written as a sample-weighted regression loss:\n",
    "\\begin{align}\n",
    "E\\left[ \\tilde{D}^2 \\left(\\tilde{Y}/\\tilde{D} - \\tau(Z) \\right)^2 \\right]\n",
    "\\end{align}\n",
    "We are predicting $\\tilde{Y}/\\tilde{D}$ form $Z$, using sample weights $\\tilde{D}^2$.\n",
    "\n",
    "If the true CATE $\\theta(X)$ did not only depend on $Z$, but on the bigger set of covariates $X$, then this method estimates the best projection of the CATE $\\theta(X)$ on the space of functions of $Z$, albeit, in a weighted manner, weighted by the variance of the treatment. More formally, the method estimates the minimizer of the following loss:\n",
    "\\begin{align}\n",
    "\\min_{\\tau} E\\left[\\left(\\theta(X) - \\tau(Z)\\right)^2 \\, \\text{Var}(D\\mid X)\\right]\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4Xe2bn2R3DBq"
   },
   "outputs": [],
   "source": [
    "# rlearner\n",
    "yres = y - res_preds\n",
    "Dres = D - prop_preds\n",
    "Dres = np.clip(Dres, 1e-6, np.inf) * (Dres >= 0) + np.clip(Dres, -np.inf, -1e-6) * (Dres < 0)\n",
    "\n",
    "rlearner = auto_weighted_reg().fit(Z.values, yres / Dres, sample_weight=Dres**2)\n",
    "\n",
    "rlearner_cates = rlearner.predict(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5fszlKn23DBr"
   },
   "source": [
    "## Constant Effect DR-Learner\n",
    "\n",
    "We also add a heavily regularized CATE model that predicts the ATE using the doubly robust pseudo outcomes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tInYN9B43DBr"
   },
   "outputs": [],
   "source": [
    "drlearner_const = make_pipeline(PolynomialFeatures(degree=0, include_bias=True),\n",
    "                                LinearRegression(fit_intercept=False))\n",
    "drlearner_const.fit(Z, dr_preds)\n",
    "drlearner_const_cates = drlearner_const.predict(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kL_PQEzr3DBr"
   },
   "source": [
    "# Interpretability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kBN_jArh3DBr"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 7))\n",
    "tree = DecisionTreeRegressor(max_depth=2, min_samples_leaf=100,\n",
    "                             min_impurity_decrease=.0001, random_state=random_seed)\n",
    "tree.fit(Ztest, drlearner.predict(Ztest))\n",
    "plot_tree(tree, filled=True, feature_names=list(Ztest.columns), fontsize=14, impurity=False)\n",
    "plt.savefig(f'{data}-distill-tree-cate-dr.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GsvtTwi73DBr"
   },
   "outputs": [],
   "source": [
    "import shap\n",
    "\n",
    "subsample = np.random.choice(Z.shape[0], min(10000, Z.shape[0]), replace=False)\n",
    "explainer = shap.TreeExplainer(drlearner)\n",
    "shap_values = explainer.shap_values(Z.iloc[subsample])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fV0MtPaq3DBr"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "shap.summary_plot(shap_values, Z.iloc[subsample], show=False)\n",
    "plt.savefig(f'{data}-shap-values-dr.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ll3UieR53DBr"
   },
   "source": [
    "# Causal Score Estimation and Definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vV9gaEnK3DBr"
   },
   "source": [
    "We want to be able to select among all these different meta learners. For this reason we will use scoring functions that can evaluate the performance of an arbitrary CATE function and is not tailored to any particular methodology. This way we can evaluate all methods using the same score on the validation set and select the best among the methods, or ensemble the methods using this scoring metric. We will describe one such meta scores, the `DR-score`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MOF5q1eW3DBr"
   },
   "source": [
    "## The Doubly Robust Score (DR-Score)\n",
    "\n",
    "The doubly robust score calculates the doubly robust proxy variables $Y_i^{DR}(g,p)$, where $g,p$ are fitted on the training set (or on the validatino in a cross-fitting manner). Then for any candidate CATE model $\\hat{\\tau}$, the DR-score is the decrease in residual variance of $Y^{DR}(g,p)$ based on the candidate model as compared to a constant effect model.\n",
    "\\begin{align}\n",
    "\\text{DRscore}(\\hat{\\tau}) := E_n\\left[\\left(Y^{DR}(g,p) - \\tau_{\\text{constant}}(Z)\\right)^2\\right] - E_n\\left[\\left(Y^{DR}(g,p) - \\hat{\\tau}(Z)\\right)^2\\right]\n",
    "\\end{align}\n",
    "We can also consider normalizing by the residual variance of the constant effect model\n",
    "\\begin{align}\n",
    "\\text{DRscore}(\\hat{\\tau}) := \\frac{E_n\\left[\\left(Y^{DR}(g,p) - \\tau_{\\text{constant}}(Z)\\right)^2\\right] - E_n\\left[\\left(Y^{DR}(g,p) - \\hat{\\tau}(Z)\\right)^2\\right]}{E_n\\left[\\left(Y^{DR}(g,p) - \\tau_{\\text{constant}}(Z)\\right)^2\\right]}\n",
    "\\end{align}\n",
    "for more interpretability, as this last quanity lies in $[-\\infty, 1]$, and for any decent model it should be a positive number in $[0,1]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "51g9jn1r3DBr"
   },
   "outputs": [],
   "source": [
    "reg_zero_on_train = modely().fit(X[D == 0], y[D == 0])\n",
    "reg_one_on_train = modely().fit(X[D == 1], y[D == 1])\n",
    "prop_on_train = auto_clf().fit(X, D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zkDeQsGf3DBr"
   },
   "outputs": [],
   "source": [
    "def calculate_dr_outcomes(reg_zero, reg_one, prop, Xval, Dval, yval):\n",
    "\n",
    "    reg_zero_preds_t = reg_zero.predict(Xval)\n",
    "    reg_one_preds_t = reg_one.predict(Xval)\n",
    "    reg_preds_t = reg_zero_preds_t * (1 - Dval) + reg_one_preds_t * Dval\n",
    "    prop_preds = prop.predict(Xval)\n",
    "\n",
    "    dr = reg_one_preds_t - reg_zero_preds_t\n",
    "    reisz = (Dval - prop_preds) / np.clip(prop_preds * (1 - prop_preds), .01, np.inf)\n",
    "    dr += (yval - reg_preds_t) * reisz\n",
    "\n",
    "    return dr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UNO4sJhG3DBs"
   },
   "outputs": [],
   "source": [
    "dr_val = calculate_dr_outcomes(reg_zero_on_train, reg_one_on_train, prop_on_train,\n",
    "                               Xval, Dval, yval)\n",
    "\n",
    "overall_ate_val_dr = np.mean(dr_preds)\n",
    "\n",
    "\n",
    "def drscore(cate_preds):\n",
    "    drscore_t = np.mean((dr_val - cate_preds)**2)\n",
    "    drscore_b = np.mean((dr_val - overall_ate_val_dr)**2)\n",
    "    return drscore_b - drscore_t\n",
    "\n",
    "\n",
    "def drscore_delta(cate_preds1, cate_preds2=overall_ate_val_dr):\n",
    "    drscore_t = np.mean((dr_val - cate_preds1)**2)\n",
    "    drscore_b = np.mean((dr_val - cate_preds2)**2)\n",
    "    stderr = np.std((dr_val - cate_preds2)**2 - (dr_val - cate_preds1)**2) / np.sqrt(cate_preds1.shape[0])\n",
    "    return drscore_b - drscore_t, stderr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ppqpu7Ai3DBs"
   },
   "source": [
    "## Compare models with confidence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EGFxBkP13DBs"
   },
   "outputs": [],
   "source": [
    "names = ['slearner', 'tlearner', 'xlearner', 'drlearner', 'rlearner']\n",
    "models = [slearner, tlearner, xlearner, drlearner, rlearner]\n",
    "\n",
    "points, ses = zip(*[drscore_delta(model.predict(Zval)) for model in models])\n",
    "points = np.array(points)\n",
    "ses = np.array(ses)\n",
    "\n",
    "plt.errorbar(names, points, yerr=[1.96 * ses, 1.96 * ses])\n",
    "plt.title('Model vs Constant ATE Baseline')\n",
    "plt.savefig(f'{data}-model-comparison.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W1KNtEzA3DBs"
   },
   "source": [
    "## Score CATE Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1VWaerd83DBs"
   },
   "outputs": [],
   "source": [
    "scorer = drscore\n",
    "score_name = 'DRscore'\n",
    "names = ['slearner', 'tlearner', 'xlearner', 'drlearner', 'rlearner', 'drlearner_const']\n",
    "models = [slearner, tlearner, xlearner, drlearner, rlearner, drlearner_const]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sA08G8E73DBs"
   },
   "outputs": [],
   "source": [
    "scores = [scorer(model.predict(Zval)) for model in models]\n",
    "print([f'{name}: {score:.4f}' for name, score in zip(names, scores)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "K1iLxMU53DBs"
   },
   "outputs": [],
   "source": [
    "for name, model in zip(names, models):\n",
    "    score = drscore_delta(model.predict(Zval))\n",
    "    print(f'{name}: {score[0]:.4f} ({score[1]:.4f}) '\n",
    "          f'[{score[0] - 1.96 * score[1]:.4f}, {score[0] + 1.96 * score[1]:.4f}]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j4pdLvQn3DBs"
   },
   "source": [
    "# Plotting CATE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3gVCgVXG3DBs"
   },
   "outputs": [],
   "source": [
    "names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zc8rfewA3DBs"
   },
   "outputs": [],
   "source": [
    "fnames = ['Single Learner', 'Two Learner', 'Cross Learner',\n",
    "          'Doubly Robust Learner', 'Residual Learner', 'Doubly Robust ATE']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8l3Oob2B3DBs"
   },
   "outputs": [],
   "source": [
    "if plot:\n",
    "    grid = np.unique(np.percentile(X[xfeat], np.arange(0, 102, 2)))\n",
    "    Xpd = pd.DataFrame(np.tile(np.median(X, axis=0, keepdims=True), (len(grid), 1)),\n",
    "                       columns=X.columns)\n",
    "    Xpd[xfeat] = grid\n",
    "    plt.figure(figsize=(15, 15))\n",
    "    for it, (fname, name, model, score) in enumerate(zip(fnames, names, models, scores)):\n",
    "        plt.subplot(3, len(models) // 3, it + 1)\n",
    "        preds = model.predict(Xpd[hetero_feats])\n",
    "        plt.plot(Xpd[xfeat], preds, linewidth=5, color='green', label=name)\n",
    "        plt.plot(Xpd[xfeat], pred_df['mean'], '--',\n",
    "                 color='orange', label='$Y(\\\\hat{\\\\eta})\\\\sim ' + f'{blp_formula_short}$')\n",
    "        plt.fill_between(Xpd[xfeat], pred_df['mean_ci_lower'], pred_df['mean_ci_upper'], color='orange', alpha=.1)\n",
    "        plt.plot(Xpd[xfeat], pred_df2['mean'], '--',\n",
    "                 color='blue', label='$Y(\\\\hat{\\\\eta})\\\\sim ' + f'{blp_formula2_short}$')\n",
    "        plt.fill_between(Xpd[xfeat], pred_df2['mean_ci_lower'], pred_df2['mean_ci_upper'], color='blue', alpha=.1)\n",
    "        if semi_synth:\n",
    "            plt.plot(Xpd[xfeat], true_cate(Xpd), label='True')\n",
    "            plt.title(f'{score_name}={score:.4f}, True RMSE={rmse(true_cate(X), model.predict(Z)):.5f}')\n",
    "        else:\n",
    "            plt.title(f'{fname}. ATE={np.mean(model.predict(Zval)):.0f}', fontsize=14)\n",
    "        plt.legend()\n",
    "        plt.xlabel(xfeat, fontsize=14)\n",
    "        plt.ylabel('CATE', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(f'{data}-metalearners.png', dpi=600)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "01LQqjmT3DBt"
   },
   "source": [
    "# Causal Model Selection and Ensembling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4rxFluNS3DBt"
   },
   "source": [
    "We can also use these scores to create an ensemble CATE model of the different methods based on the score performance. We want to create a new CATE model $\\tau_E$ that is a weighted linear combination of all the cate models, i.e.\n",
    "\\begin{align}\n",
    "\\tau_E(Z) = \\bar{\\tau} + \\sum_{m=1}^M w_m \\tau_m(Z)\n",
    "\\end{align}\n",
    "One way to achieve that is to choose the model with the best score. However, the discontinuity in this selection process can be statistically suboptimal. One way to achieve this in a smoother manner is to perform Stacking to construct the ensemble weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sgpxqIRN3DBt"
   },
   "outputs": [],
   "source": [
    "from sklearn.base import BaseEstimator\n",
    "\n",
    "\n",
    "class Ensemble(BaseEstimator):\n",
    "\n",
    "    def __init__(self, names, models, weights, intercept=0):\n",
    "        self.names = names\n",
    "        self.models = models\n",
    "        self.weights = weights\n",
    "        self.intercept = intercept\n",
    "\n",
    "    def predict(self, X):\n",
    "        wcate = np.sum(self.weights.reshape((-1, 1)) * np.array([m.predict(X) for m in self.models]), axis=0)\n",
    "        return self.intercept + wcate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ErrJdAlL3DBt"
   },
   "source": [
    "We perform Stacking by fitting a (potentially $\\ell_1$-penalized) linear model to minimize the corresponding loss. For the case of the DR score this boils down to a penalized linear regression, predicting $Y^{DR}(g,p)$ on the validation set, using $\\tau_1(Z),\\ldots, \\tau_M(Z)$ as technical regressors:\n",
    "\\begin{align}\n",
    "\\min_{w} E_n\\left[ \\left(Y^{DR}(g,p) - \\bar{\\tau} - \\sum_{m=1}^M w_m \\tilde{\\tau}_m(Z)\\right)^2 \\right] + \\lambda \\text{Penalty}(w)\n",
    "\\end{align}\n",
    "where $\\bar{\\tau}$ is the ATE estimate based on the doubly robust method on the training data and $\\tilde{\\tau}_m(Z)$ are the de-meaned CATE predictions from each model, i.e. $\\tilde{\\tau}_m(Z) = \\tau_m(Z) - E_n[\\tau_m(Z)]$ where the empirical expectation is over the training data.\n",
    "\n",
    "This way we are constructing an ensemble of the CATE on top of the best constant prediction on the training data and penalizing only the offset from this ATE. We avoid adding a coefficient around $\\bar{\\tau}$ to avoid introducing noise to the ATE estimate, due to the smaller sample in the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WoLIOEBs3DBt"
   },
   "outputs": [],
   "source": [
    "from sklearn.linear_model import RidgeCV, LinearRegression\n",
    "\n",
    "F = np.array([m.predict(Zval) for m in models]).T\n",
    "meansF = np.mean(F, axis=0)\n",
    "F = F - meansF\n",
    "\n",
    "stack_models = {}\n",
    "stack_variants = [('OLS', LinearRegression(fit_intercept=False)),\n",
    "                  ('OLS (positive)', LinearRegression(fit_intercept=False, positive=True)),\n",
    "                  ('LassoCV', LassoCV(fit_intercept=False)),\n",
    "                  ('LassoCV (positive)', LassoCV(fit_intercept=False, positive=True)),\n",
    "                  ('RidgeCV', RidgeCV(fit_intercept=False))]\n",
    "\n",
    "for stackname, stacker in stack_variants:\n",
    "    stacker.fit(F, dr_val - np.mean(dr_preds))\n",
    "    intercept = np.mean(dr_preds) - meansF @ stacker.coef_\n",
    "    stack_models[stackname] = Ensemble(names, models, stacker.coef_, intercept)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J3fVPlhJzDCi"
   },
   "source": [
    "Q-aggregation optimizes a penalized objective of the form:\n",
    "\\begin{align}\n",
    "\\min_{w\\in \\Delta(M)} (1 - \\nu) E_n\\left[ \\left(Y^{DR}(g,p) - \\bar{\\tau} - \\sum_{m=1}^M w_m \\tilde{\\tau}_m(Z)\\right)^2 \\right] + \\nu \\sum_{m=1}^M w_m E_n\\left[ \\left(Y^{DR}(g,p) - \\bar{\\tau} - \\tilde{\\tau}_m(Z)\\right)^2 \\right]\n",
    "\\end{align}\n",
    "where $\\Delta(M)$ is the simplex in $M$ dimensions, i.e.\n",
    "\\begin{align}\n",
    "\\Delta(M) := \\{w\\in R^M: w_i\\geq 0, \\sum_{m=1}^M w_m=1\\}.\n",
    "\\end{align}\n",
    "\n",
    "In other words, it chooses a convex combination of models, but penalizes models which do not perform well individually. Note that when $\\nu=0$, then this boils down to what is known as \"convex regression\", i.e. OLS regression subject to a simplex constraint on the coefficients. When $\\nu=1$, then the above boils down to choosing the single best performing model. Hence, Q-aggregation interpolates between the best convex combination of models and the single best performing model. It enjoys favorable theoretical guarantees as compared to either of the two extremes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xo2fMdpM3DBt"
   },
   "outputs": [],
   "source": [
    "def instance(nu, U, y):\n",
    "    n = y.shape[0]\n",
    "    ploss = np.mean((y.reshape(1, -1) - U)**2, axis=1)\n",
    "\n",
    "    def loss(x):\n",
    "        return np.mean((y - U.T @ x)**2)\n",
    "\n",
    "    def qfunction(x):\n",
    "        return (1 - nu) * loss(x) + nu * x @ ploss\n",
    "\n",
    "    def grad_q(x):\n",
    "        return - 2 * (1 - nu) * U @ (y - U.T @ x) / n + nu * ploss\n",
    "\n",
    "    return loss, qfunction, grad_q, ploss\n",
    "\n",
    "\n",
    "def opt(K, qfunction, grad_q):\n",
    "    res = scipy.optimize.minimize(qfunction, np.ones(K) / K, jac=grad_q, bounds=[(0, 1)] * K,\n",
    "                                  constraints=scipy.optimize.LinearConstraint(np.ones((1, K)), lb=1, ub=1),\n",
    "                                  tol=1e-18)\n",
    "    return res.x\n",
    "\n",
    "\n",
    "def qagg(F, y, nu=.5):\n",
    "    scale = max(np.max(np.abs(F)), np.max(np.abs(y)))\n",
    "    loss, qfunction, grad_q, ploss = instance(nu, F.T / scale, y / scale)\n",
    "    return opt(F.shape[1], qfunction, grad_q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XSpQhhkG3DBt"
   },
   "outputs": [],
   "source": [
    "F = np.array([m.predict(Zval) for m in models]).T\n",
    "weights = qagg(F, dr_val)\n",
    "stack_models['Q-aggregation'] = Ensemble(names, models, weights, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "enVDiqSk3DBt"
   },
   "outputs": [],
   "source": [
    "F = np.array([m.predict(Zval) for m in models]).T\n",
    "weights = qagg(F, dr_val, nu=0.0)\n",
    "stack_models['OLS (simplex)'] = Ensemble(names, models, weights, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2LYwKVe03DBt"
   },
   "outputs": [],
   "source": [
    "F = np.array([m.predict(Zval) for m in models]).T\n",
    "weights = qagg(F, dr_val, nu=1.0)\n",
    "stack_models['Best'] = Ensemble(names, models, weights, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TWZucUU53DBt"
   },
   "outputs": [],
   "source": [
    "plt.plot(names, stack_models['Q-aggregation'].weights, label='Q-aggregation')\n",
    "plt.plot(names, stack_models['Best'].weights, label='Best')\n",
    "plt.plot(names, stack_models['OLS (simplex)'].weights, label='OLS (simplex)')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pfJ7cdkA3DBt"
   },
   "outputs": [],
   "source": [
    "if plot:\n",
    "    stack_ensemble = stack_models['Q-aggregation']\n",
    "    grid = np.unique(np.percentile(X[xfeat], np.arange(0, 102, 2)))\n",
    "    Xpd = pd.DataFrame(np.tile(np.median(X, axis=0, keepdims=True), (len(grid), 1)),\n",
    "                       columns=X.columns)\n",
    "    Xpd[xfeat] = grid\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    plt.plot(Xpd[xfeat], stack_ensemble.predict(Xpd[hetero_feats]),\n",
    "             linewidth=5, color='green',\n",
    "             label='Q-aggregate')\n",
    "    plt.plot(Xpd[xfeat], pred_df['mean'], '--', color='orange',\n",
    "             label='$Y(\\\\hat{\\\\eta})\\\\sim ' + f'{blp_formula_short}$')\n",
    "    plt.fill_between(Xpd[xfeat], pred_df['mean_ci_lower'], pred_df['mean_ci_upper'], color='orange', alpha=.1)\n",
    "    plt.plot(Xpd[xfeat], pred_df2['mean'], '--', color='blue',\n",
    "             label='$Y(\\\\hat{\\\\eta})\\\\sim ' + f'{blp_formula2_short}$')\n",
    "    plt.fill_between(Xpd[xfeat], pred_df2['mean_ci_lower'], pred_df2['mean_ci_upper'], color='blue', alpha=.1)\n",
    "    if semi_synth:\n",
    "        plt.plot(Xpd[xfeat], true_cate(Xpd), '--', label='True')\n",
    "        plt.title(f'{score_name}={scorer(stack_ensemble.predict(Zval)):.5f}, '\n",
    "                  f'True RMSE={rmse(true_cate(X), stack_ensemble.predict(Z)):.5f}')\n",
    "    else:\n",
    "        plt.title('Q-aggregate CATE Model, '\n",
    "                  f'ATE={np.mean(stack_ensemble.predict(Zval)):.5f}')\n",
    "    plt.xlabel(xfeat)\n",
    "    plt.ylabel('CATE')\n",
    "    plt.legend()\n",
    "    plt.savefig(f'{data}-qagg.png', dpi=600)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vbuKWzLY3DBt"
   },
   "outputs": [],
   "source": [
    "if plot:\n",
    "    grid = np.unique(np.percentile(X[xfeat], np.arange(0, 102, 2)))\n",
    "    Xpd = pd.DataFrame(np.tile(np.median(X, axis=0, keepdims=True), (len(grid), 1)),\n",
    "                       columns=X.columns)\n",
    "    Xpd[xfeat] = grid\n",
    "    plt.figure(figsize=(15, 15))\n",
    "    for it, (name, model) in enumerate(stack_models.items()):\n",
    "        plt.subplot(4, len(stack_models) // 4, it + 1)\n",
    "        preds = model.predict(Xpd[hetero_feats])\n",
    "        plt.plot(Xpd[xfeat], preds, linewidth=5, color='green', label=name)\n",
    "        plt.plot(Xpd[xfeat], pred_df['mean'], '--', color='orange',\n",
    "                 label='$Y(\\\\hat{\\\\eta})\\\\sim ' + f'{blp_formula_short}$')\n",
    "        plt.fill_between(Xpd[xfeat], pred_df['mean_ci_lower'], pred_df['mean_ci_upper'],\n",
    "                         color='orange', alpha=.1)\n",
    "        plt.plot(Xpd[xfeat], pred_df2['mean'], '--', color='blue',\n",
    "                 label='$Y(\\\\hat{\\\\eta})\\\\sim ' + f'{blp_formula2_short}$')\n",
    "        plt.fill_between(Xpd[xfeat], pred_df2['mean_ci_lower'], pred_df2['mean_ci_upper'],\n",
    "                         color='blue', alpha=.1)\n",
    "        if semi_synth:\n",
    "            plt.plot(Xpd[xfeat], true_cate(Xpd), label='True')\n",
    "            plt.title(f'True RMSE={rmse(true_cate(X), model.predict(Z)):.5f}')\n",
    "        else:\n",
    "            plt.title(f'{name}. ATE={np.mean(model.predict(Zval)):.0f}', fontsize=14)\n",
    "        plt.legend()\n",
    "        plt.xlabel(xfeat, fontsize=14)\n",
    "        plt.ylabel('CATE', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(f'{data}-stacked_ensembles.png', dpi=600)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "epqDWStA3DBt"
   },
   "source": [
    "We use one of these ensembles as our final best model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6kiuJCDq3DBu"
   },
   "outputs": [],
   "source": [
    "overall_best = stack_models['Q-aggregation']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wPrRJ_YJ3DBu"
   },
   "source": [
    "# Interpret with SHAP Values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IqLZFscU3DBu"
   },
   "outputs": [],
   "source": [
    "import shap\n",
    "\n",
    "subsample = np.random.choice(Z.shape[0], min(1000, Z.shape[0]), replace=False)\n",
    "explainer = shap.Explainer(overall_best.predict, masker=Z.iloc[subsample])\n",
    "subsample = np.random.choice(Z.shape[0], min(100, Z.shape[0]), replace=False)\n",
    "shap_values = explainer(Z.iloc[subsample])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RCKhJQap3DBu"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "shap.summary_plot(shap_values, Z.iloc[subsample], show=False)\n",
    "plt.savefig(f'{data}-shap-values.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MCQ5SWxB3DBu"
   },
   "source": [
    "# Interpret with Distillation Trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j7P2FaBV3DBu"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 7))\n",
    "tree = DecisionTreeRegressor(max_depth=2, min_samples_leaf=100,\n",
    "                             min_impurity_decrease=.00001, random_state=random_seed)\n",
    "tree.fit(Ztest, overall_best.predict(Ztest))\n",
    "plot_tree(tree, filled=True, feature_names=list(Ztest.columns), fontsize=14, impurity=False)\n",
    "plt.savefig(f'{data}-distill-tree-cate.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nmOkLCe43DBu"
   },
   "source": [
    "# Validation Tests on Test Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b1iDNGgr3DBu"
   },
   "source": [
    "Now that we have a selected a winning CATE model (or ensemble), we can run a set of hypothesis tests and other diagnostic metrics on the test set, to see if the model really picked up some dimensions of effect heterogeneity and satisfies some self-conistency checks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vJQydwCM3DBu"
   },
   "source": [
    "### Hypothesis Test Based on Doubly Robust Best-Linear Predictor of CATE using model of CATE\n",
    "\n",
    "If we calculate the doubly robust pseudo outcomes $Y^{DR}(g,p)$ on the test set, using cross-fitting to estimate the models $g,p$, then we know that if the model of the CATE $\\hat{\\tau}$ is good, then the best linear predictor of the true CATE using $(1, \\hat{\\tau}(Z))$ as features should have a statistically significant coefficient on the feature associated with the CATE model. In fact, in an ideal world it should have a coefficient of $1$. Thus we can run such a significance test to measure whether the CATE model $\\tau_m$ has picked up anything signal that is correlated with the true CATE. Note that if $\\theta(X)$ is the true CATE $E[Y(1)-Y(0)\\mid X]$, then the coefficient associated with $\\hat{\\tau}$ in this regression can be interpreted as identifying the quantity:\n",
    "\\begin{align}\n",
    "\\beta_1 := \\frac{\\text{Cov}(\\theta(X), \\hat{\\tau}(Z))}{\\text{Var}(\\hat{\\tau}(Z))} = \\frac{\\text{Cov}(Y(1)-Y(0), \\hat{\\tau}(Z))}{\\text{Var}(\\hat{\\tau}(Z))}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O8nCBkmJ3DBu"
   },
   "outputs": [],
   "source": [
    "dr_test = calculate_dr_outcomes(reg_zero_on_train, reg_one_on_train, prop_on_train,\n",
    "                                Xtest, Dtest, ytest)\n",
    "cate_test = overall_best.predict(Ztest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W9rR11483DBu"
   },
   "outputs": [],
   "source": [
    "res = OLS(dr_test, np.stack((np.ones(len(dr_test)), cate_test - np.mean(cate_test)), axis=-1))\n",
    "res = res.fit(cov_type='HC1').summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z-cG7Sdn3DBu"
   },
   "outputs": [],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wmsbXn8x3DBu"
   },
   "source": [
    "### Validation Based on Calibration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DTUlS5sw3DBu"
   },
   "source": [
    "We can measure whether each group defined by the quartile levels of CATE predictions is consistent with the out-of-sample Group ATE (GATE) for the corresponding group based on the doubly robust GATE estimate. (standard errors here ignore cluster/group correlations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UcWhw8iR3DBu"
   },
   "outputs": [],
   "source": [
    "cate_val = overall_best.predict(Zval)\n",
    "qs = np.percentile(cate_val, np.arange(0, 101, 25))\n",
    "\n",
    "gate, gate_std, group_prob = np.zeros(len(qs) - 1), np.zeros(len(qs) - 1), np.zeros(len(qs) - 1)\n",
    "predicted_gate = np.zeros(len(qs) - 1)\n",
    "for it in range(len(qs) - 1):\n",
    "    # samples in the [q[it], q[it+1]) quantile group of predicted CATEs\n",
    "    inds = (qs[it] <= cate_test) & (cate_test <= qs[it + 1])\n",
    "    gate[it] = np.mean(dr_test[inds])  # DR estimate of group average treatment effect (GATE)\n",
    "    gate_std[it] = np.std(dr_test[inds]) / np.sqrt(np.sum(inds))  # standard error of GATE\n",
    "    group_prob[it] = np.mean(inds)  # probability mass of group\n",
    "    predicted_gate[it] = np.mean(cate_test[inds])  # GATE as calculated from CATE model\n",
    "\n",
    "# weighted average calibration error of cate model\n",
    "cal = np.sum(group_prob * np.abs(gate - predicted_gate))\n",
    "# weighted average calibration error of a constant cate model\n",
    "calbase = np.sum(group_prob * np.abs(gate - np.mean(dr_test)))\n",
    "# calibration score\n",
    "calscore = 1 - cal / calbase\n",
    "plt.title(f'CalScore={calscore:.4f}')\n",
    "plt.errorbar(predicted_gate, gate, yerr=1.96 * gate_std, fmt='o')\n",
    "plt.xlabel('Predicted GATE based on CATE model')\n",
    "plt.ylabel('Doubly Robust GATE estimate')\n",
    "plt.savefig(f'{data}-calibration-score.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pzCkkJRJ3DBu"
   },
   "source": [
    "We can also try to interpret what are the differences of characteristics between the top and bottom CATE groups; if we find that they have statistically significantly different GATEs. We can do that by either reporting the mean values of the covariates in the two groups or building some interpretable classification model that distinguishes between the two groups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IaBAYLhk3DBv"
   },
   "outputs": [],
   "source": [
    "group1 = (qs[0] <= cate_test) & (cate_test < qs[3])\n",
    "group2 = (qs[-2] <= cate_test) & (cate_test < qs[-1])\n",
    "Ztest1 = Ztest[group1]\n",
    "Ztest2 = Ztest[group2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "06zRJgMF3DBv"
   },
   "outputs": [],
   "source": [
    "diff_stderr = np.std(Ztest1, axis=0) / np.sqrt(Ztest1.shape[0]) + np.std(Ztest2, axis=0) / np.sqrt(Ztest2.shape[0])\n",
    "df = pd.DataFrame({'group1 means': np.mean(Ztest1, axis=0),\n",
    "                   'group1 s.e.': np.std(Ztest1, axis=0) / np.sqrt(Ztest1.shape[0]),\n",
    "                   'group2 means': np.mean(Ztest2, axis=0),\n",
    "                   'group2 s.e.': np.std(Ztest2, axis=0) / np.sqrt(Ztest2.shape[0]),\n",
    "                   'group1 means - group2 means': np.mean(Ztest1, axis=0) - np.mean(Ztest2, axis=0),\n",
    "                   'diff s.e.': diff_stderr})\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q2C80imW3DBv"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 7))\n",
    "tree = DecisionTreeClassifier(max_depth=2, min_samples_leaf=50, min_impurity_decrease=.01)\n",
    "tree.fit(pd.concat((Ztest1, Ztest2)),\n",
    "         np.concatenate((np.zeros(len(Ztest1)), np.ones(len(Ztest2)))))\n",
    "plot_tree(tree, filled=True, feature_names=list(Ztest1.columns),\n",
    "          class_names=['group1', 'group2'], fontsize=14, impurity=False)\n",
    "plt.savefig(f'{data}-distill-groups-tree-cate.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "90lv68By3DBv"
   },
   "source": [
    "### Validation Based on Uplift Curves"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O60J05PF3DBv"
   },
   "source": [
    "These curves are related to \"prioritization\" or \"stratification\" implications of the CATE model. What if we target to treat a $q$-percentage of the population. Then if we trust and follow the CATE model, then we should be offering the treatment to the parts of the population that the CATE model predicts have a CATE larger than the $1-q$-th percentile of the CATE distribution as produced by the CATE.\n",
    "\n",
    "In this case, we might care about the group average treatment effect of the group of people that are treated under such a prioritization rule and how that compares with the overall average treatment effect. Ideally, the group average treatment effect should be larger than the overall average treatment effect if the prioritization rule is correct.\n",
    "\n",
    "Thus if we have a rough estimate $\\hat{\\mu}(q)$ of the $1-q$-quantile $\\mu(q)$ of the CATE predictions distribution we can calculate:\n",
    "\\begin{align}\n",
    "\\tau(q) =~& E[Y(1) - Y(0) \\mid \\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)] - E[Y(1) - Y(0)]\\\\\n",
    "=~& E\\left[(Y(1) - Y(0)) \\frac{1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}}{\\Pr(\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q))}\\right] - E[Y(1) - Y(0)]\n",
    "~=~ \\text{Cov}\\left(Y(1) - Y(0), \\frac{1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}}{\\Pr(\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q))}\\right)\n",
    "\\end{align}\n",
    "We can get such good quantile estimates out-of-sample, by looking at the $1-q$-the empirical quantile of the cate predictions on a sample other than the one that we are using to calculate $\\tau(q)$. For instance, we can use the training/validation samples but not the test samples. This gives us a set of thresholds $\\hat{\\mu}(q_1), \\ldots, \\hat{\\mu}(q_m)$ that are good approximations of the quantile of the predicted CATE distribution. If we target to treat every person with CATE greater than $\\hat{\\mu}(q_t)$ we should expect roughly $q_t$ fraction of the population to be treated.\n",
    "\n",
    "This quantity is identified by replacing the individual effects with the doubly robust pseudo-outcomes:\n",
    "\\begin{align}\n",
    "\\tau(q) =~& E\\left[Y^{DR}(g,p) \\mid \\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\right] - E\\left[Y^{DR}(g,p)\\right]\n",
    "~=~ \\text{Cov}\\left(Y^{DR}(g,p), \\frac{1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}}{\\Pr(\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q))}\\right)\n",
    "\\end{align}\n",
    "\n",
    "We can then plot the curve $\\tau(q)$, understand the Area Under the Curve:\n",
    "\\begin{align}\n",
    "AUTOC := \\int_0^1 \\tau(q) dq\n",
    "\\end{align}\n",
    "and run tests whether that area is positive or that there is any positive point in the curve. These are tests that indicate that the CATE model detected heterogeneity with statistical significance. Moreover, the larger the Area Under the Curve, the better the CATE model is at treatment prioritization. (standard errors here ignore cluster/group correlations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "91QWG6jy3DBv"
   },
   "source": [
    "We also need to be able to handle tie-breaking in case where the $\\hat{\\mu}(q)$ curve contains regions where the quantile is flat. We will use the following type of formula for that case\n",
    "\\begin{align}\n",
    "\\tau(q) =~& E[A \\mid T<C, (T=C, U<q)] - E[A]\\\\\n",
    "=~& E\\left[A \\frac{1\\{T<C, (T=C, U<q)\\}}{Pr}\\right] - E[A]\\\\\n",
    "=~& E\\left[A \\frac{1\\{T<C\\} + 1\\{T=C\\} q}{Pr(T<C) + Pr(T=C) q}\\right] - E[A]\\\\\n",
    "=~& Cov\\left(A, \\frac{1\\{T<C\\} + 1\\{T=C\\} q}{Pr(T<C) + Pr(T=C) q}\\right)\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gnqytmgF3DBv"
   },
   "outputs": [],
   "source": [
    "# Based on out-of-sample CATE thresholds\n",
    "ugrid = np.linspace(5, 95, 50)\n",
    "qs = np.percentile(overall_best.predict(Zval), ugrid)\n",
    "\n",
    "toc, toc_std, group_prob = np.zeros(len(qs)), np.zeros(len(qs)), np.zeros(len(qs))\n",
    "true_toc = np.zeros(len(qs))\n",
    "toc_psi = np.zeros((len(qs), dr_test.shape[0]))  # influence function representation of the TOC at each quantile\n",
    "n = len(dr_test)\n",
    "ate = np.mean(dr_test)\n",
    "uqs, ind, cnt = np.unique(qs, return_counts=True, return_index=True)\n",
    "for it, u in enumerate(ugrid):\n",
    "    # calculating tie breaking fraction\n",
    "    intv = np.searchsorted(ugrid[ind], u, side='right') - 1\n",
    "    q, i, c = uqs[intv], ind[intv], cnt[intv]\n",
    "    nextu = (ugrid[i + c + 1] if i + c + 1 < len(ugrid) else 100)\n",
    "    remnant = nextu - u\n",
    "    frac = remnant / (nextu - ugrid[i])\n",
    "\n",
    "    inds = (q < cate_test) + (cate_test == q) * frac  # group with larger CATE prediction than the q-th quantile\n",
    "    group_prob = np.mean(inds)\n",
    "    toc[it] = np.mean(dr_test * inds / group_prob) - ate  # tau(q) := E[Y(1) - Y(0) | tau(X) >= q[it]] - E[Y(1) - Y(0)]\n",
    "    # influence function for the tau(q); it is a standard influence function of a \"covariance\"\n",
    "    toc_psi[it, :] = (dr_test - ate) * (inds / group_prob - 1) - toc[it]\n",
    "    toc_std[it] = np.sqrt(np.mean(toc_psi[it]**2) / n)  # standard error of tau(q)\n",
    "    if semi_synth:\n",
    "        true_toc[it] = np.mean((true_cate(Xtest) - np.mean(true_cate(Xtest))) * (inds / group_prob - 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hDgc8Mvp3DBv"
   },
   "outputs": [],
   "source": [
    "plt.errorbar(100 - ugrid, toc, yerr=1.96 * toc_std, fmt='o', label='Est. TOC')\n",
    "plt.plot(100 - ugrid, np.zeros(len(ugrid)))\n",
    "if semi_synth:\n",
    "    plt.plot(100 - ugrid, true_toc, 'o', label='True TOC')\n",
    "plt.xlabel(\"Percentage treated\")\n",
    "plt.ylabel(\"Gain in Average Effect of Treated by CATE over Random\")\n",
    "plt.legend()\n",
    "plt.savefig(f'{data}-toc.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oF0jSb2e3DBv"
   },
   "outputs": [],
   "source": [
    "V = toc_psi @ toc_psi.T / toc_psi.shape[1]\n",
    "sigmas = np.sqrt(np.diag(V))\n",
    "V /= sigmas.reshape(-1, 1)\n",
    "V /= sigmas.reshape(1, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5rKM_8rw3DBv"
   },
   "outputs": [],
   "source": [
    "epsilon = np.random.multivariate_normal(np.zeros(V.shape[0]), V, size=(100000))\n",
    "uniform_critical_value = np.percentile(np.max(np.abs(epsilon), axis=1), 95)\n",
    "uniform_critical_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AvxtgjlZ3DBv"
   },
   "outputs": [],
   "source": [
    "plt.errorbar(100 - ugrid, toc, yerr=uniform_critical_value * toc_std, fmt='o', label='Est. TOC')\n",
    "plt.plot(100 - ugrid, np.zeros(len(ugrid)))\n",
    "if semi_synth:\n",
    "    plt.plot(100 - ugrid, true_toc, 'o', label='True TOC')\n",
    "plt.xlabel(\"Percentage treated\")\n",
    "plt.ylabel(\"Gain in Average Effect of Treated by CATE over Random\")\n",
    "plt.legend()\n",
    "plt.savefig(f'{data}-toc-uniform-band.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tgI-KBHa3DBw"
   },
   "source": [
    "Note that if there is any point that is above the zero line, with confidence, in this curve, then the CATE model $\\hat{\\tau}$ has identified heterogeneity in the effect in a statistically significant manner. To do this we need a one-sided confidence interval, as we only care that the quantities are larger than some value with high confidence. We can then calculate the critical value for a uniform one-sided confidence interval across all the points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AT4IqQB53DBw"
   },
   "outputs": [],
   "source": [
    "uniform_one_side_critical_value = np.percentile(np.max(np.abs(epsilon), axis=1), 90)\n",
    "print(uniform_one_side_critical_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yF9Qqvna3DBw"
   },
   "outputs": [],
   "source": [
    "plt.errorbar(100 - ugrid, toc,\n",
    "             yerr=[uniform_one_side_critical_value * toc_std, np.zeros(len(toc))], fmt='o', label='Est. TOC')\n",
    "plt.plot(100 - ugrid, np.zeros(len(ugrid)))\n",
    "if semi_synth:\n",
    "    plt.plot(100 - ugrid, true_toc, 'o', label='True TOC')\n",
    "plt.xlabel(\"Percentage treated\")\n",
    "plt.ylabel(\"Gain in Average Effect of Treated by CATE over Random\")\n",
    "plt.title(f'Heterogeneity Statistic: {np.max(toc - uniform_one_side_critical_value*toc_std):.4f}')\n",
    "plt.legend()\n",
    "plt.savefig(f'{data}-toc-one-sided-band.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YhUBofvZ3DBw"
   },
   "source": [
    "We can also calcualte the area under the curve and the confidence interval for that area. If the confidence interval does not contain zero, then we have again detected heterogeneity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ecJ0iz2o3DBw"
   },
   "outputs": [],
   "source": [
    "autoc_psi = np.sum(toc_psi[:-1] * np.diff(ugrid).reshape(-1, 1) / 100, 0)\n",
    "autoc = np.sum(toc[:-1] * np.diff(ugrid) / 100)\n",
    "autoc_stderr = np.sqrt(np.mean(autoc_psi**2) / n)\n",
    "print(f'AUTOC: {autoc:.4f}, s.e.: {autoc_stderr:.4f}, '\n",
    "      f'One-Sided 95% CI=[{autoc - scipy.stats.norm.ppf(.95) * autoc_stderr:.4f}, Infty]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hatM8_tn3DBw"
   },
   "source": [
    "### Qini Curve\n",
    "\n",
    "Similar to the $\\tau(q)$ curve, we can also consider a small variant of that curve that also incorporates the probability of treatment. So instead of looking just at the group average effect of the group of people that are prioritized by the CATE model and subtract the overall effect, we can look at the \"total value\" exctracted by the treated group and contrast it with the total value that would be extracted by an equally large group treated uniformly at random. This corresponds to:\n",
    "\\begin{align}\n",
    "\\tau_{\\text{QINI}}(q) :=~& \\tau(q)\\, \\Pr(\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q))\\\\\n",
    "=~& \\left(E[Y(1) - Y(0)\\mid \\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)] - E[Y(1) - Y(0)]\\right)\\, \\Pr(\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q))\\\\\n",
    "=~& E\\left[ (Y(1) - Y(0)) 1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}\\right] - E[Y(1) - Y(0)]\\, \\Pr(\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q))\\\\\n",
    "=~& E\\left[ (Y(1) - Y(0)) 1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}\\right] - E[Y(1) - Y(0)]\\, E\\left[1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}\\right]\\\\\n",
    "=~& \\text{Cov}\\left(Y(1) - Y(0), 1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}\\right)\n",
    "\\end{align}\n",
    "The latter is approximately $\\tau_{\\text{QINI}}(q) \\approx q \\tau(q)$, but might not be exactly so, due to the inaccuracy of the quantile approximation $\\hat{\\mu}(q)$.\n",
    "\n",
    "The $\\tau_{\\text{QINI}}(q)$ quantity is identified by replacing the individual effects $Y(1)-Y(0)$ with the doubly robust pseudo-outcomes:\n",
    "\\begin{align}\n",
    "\\tau_{\\text{QINI}}(q) := \\text{Cov}\\left(Y^{DR}(g,p), 1\\{\\hat{\\tau}(Z) \\geq \\hat{\\mu}(q)\\}\\right)\n",
    "\\end{align}\n",
    "\n",
    "Similarly, we can look at the area under this curve:\n",
    "\\begin{align}\n",
    "\\text{QINI} := \\int_0^1 \\tau_{\\text{QINI}}(q) dq\n",
    "\\end{align}\n",
    "which is known as the QINI coefficient. (standard errors here ignore cluster/group correlations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rinRMF573DBw"
   },
   "outputs": [],
   "source": [
    "# Based on out-of-sample CATE thresholds\n",
    "ugrid = np.linspace(5, 95, 50)\n",
    "qs = np.percentile(overall_best.predict(Zval), ugrid)\n",
    "\n",
    "toc, toc_std, group_prob = np.zeros(len(qs)), np.zeros(len(qs)), np.zeros(len(qs))\n",
    "true_toc = np.zeros(len(qs))\n",
    "toc_psi = np.zeros((len(qs), dr_test.shape[0]))\n",
    "n = len(dr_test)\n",
    "ate = np.mean(dr_test)\n",
    "uqs, ind, cnt = np.unique(qs, return_counts=True, return_index=True)\n",
    "for it, u in enumerate(ugrid):\n",
    "    # calculating tie breaking fraction\n",
    "    intv = np.searchsorted(ugrid[ind], u, side='right') - 1\n",
    "    q, i, c = uqs[intv], ind[intv], cnt[intv]\n",
    "    nextu = (ugrid[i + c + 1] if i + c + 1 < len(ugrid) else 100)\n",
    "    remnant = nextu - u\n",
    "    frac = remnant / (nextu - ugrid[i])\n",
    "\n",
    "    inds = (q < cate_test) + (cate_test == q) * frac  # group with larger CATE prediction than the q-th quantile\n",
    "    group_prob = np.mean(inds)\n",
    "    # tau(q) = q * E[Y(1) - Y(0) | tau(X) >= q[it]] - E[Y(1) - Y(0)]\n",
    "    toc[it] = np.mean((dr_test - ate) * (inds - group_prob))\n",
    "    toc_psi[it, :] = (dr_test - ate) * (inds - group_prob) - toc[it]  # influence function for the tau(q)\n",
    "    toc_std[it] = np.sqrt(np.mean(toc_psi[it]**2) / n)  # standard error of tau(q)\n",
    "    if semi_synth:\n",
    "        true_toc[it] = np.mean((true_cate(Xtest) - np.mean(true_cate(Xtest))) * (inds - group_prob))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oDFSXUDX3DBw"
   },
   "outputs": [],
   "source": [
    "plt.errorbar(100 - ugrid, toc, yerr=1.96 * toc_std, fmt='o', label='Est. QINI')\n",
    "plt.plot(100 - ugrid, np.zeros(len(ugrid)))\n",
    "if semi_synth:\n",
    "    plt.plot(100 - ugrid, true_toc, 'o', label='True QINI')\n",
    "plt.xlabel(\"Percentage treated\")\n",
    "plt.ylabel(\"Gain in Policy Value over Random Treatment\")\n",
    "plt.legend()\n",
    "plt.savefig(f'{data}-qini.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "khhtocVK3DBw"
   },
   "source": [
    "And for uniform coverage we can again use a confidence band critical value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IDjdCDcq3DBw"
   },
   "outputs": [],
   "source": [
    "V = toc_psi @ toc_psi.T / toc_psi.shape[1]\n",
    "sigmas = np.sqrt(np.diag(V))\n",
    "V /= sigmas.reshape(-1, 1)\n",
    "V /= sigmas.reshape(1, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0_RYQzxq3DBw"
   },
   "outputs": [],
   "source": [
    "epsilon = np.random.multivariate_normal(np.zeros(V.shape[0]), V, size=(100000))\n",
    "uniform_critical_value = np.percentile(np.max(np.abs(epsilon), axis=1), 95)\n",
    "uniform_critical_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g1krFyLr3DBw"
   },
   "outputs": [],
   "source": [
    "plt.errorbar(100 - ugrid, toc, yerr=uniform_critical_value * toc_std, fmt='o', label='Est. QINI')\n",
    "plt.plot(100 - ugrid, np.zeros(len(ugrid)))\n",
    "if semi_synth:\n",
    "    plt.plot(100 - ugrid, true_toc, 'o', label='True QINI')\n",
    "plt.xlabel(\"Percentage treated\")\n",
    "plt.ylabel(\"Gain in Policy Valuee over Random Treatment\")\n",
    "plt.legend()\n",
    "plt.savefig(f'{data}-qini-band.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MRp6Snf_3DBw"
   },
   "source": [
    "Or the one-sided critical value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q2lW6rQA3DBx"
   },
   "outputs": [],
   "source": [
    "uniform_one_side_critical_value = np.percentile(np.max(np.abs(epsilon), axis=1), 90)\n",
    "print(uniform_one_side_critical_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PjXVxqLM3DBx"
   },
   "outputs": [],
   "source": [
    "plt.errorbar(100 - ugrid, toc, yerr=[uniform_one_side_critical_value * toc_std, np.zeros(len(toc))],\n",
    "             fmt='o', label='Est. QINI')\n",
    "plt.plot(100 - ugrid, np.zeros(len(ugrid)))\n",
    "if semi_synth:\n",
    "    plt.plot(100 - ugrid, true_toc, 'o', label='True QINI')\n",
    "plt.xlabel(\"Percentage treated\")\n",
    "plt.ylabel(\"Gain in Average Effect of Treated by CATE over Random\")\n",
    "plt.title(f'Heterogeneity Statistic: {np.max(toc - uniform_one_side_critical_value*toc_std):.4f}')\n",
    "plt.legend()\n",
    "plt.savefig(f'{data}-qini-one-sided-band.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EFHhpI0n3DBx"
   },
   "outputs": [],
   "source": [
    "qini_psi = np.sum(toc_psi[:-1] * np.diff(ugrid).reshape(-1, 1) / 100, 0)\n",
    "qini = np.sum(toc[:-1] * np.diff(ugrid) / 100)\n",
    "qini_stderr = np.sqrt(np.mean(qini_psi**2) / n)\n",
    "print(f'QINI: {qini:.4f}, s.e.: {qini_stderr:.4f}, '\n",
    "      f'One-Sided 95% CI=[{qini - scipy.stats.norm.ppf(.95) * qini_stderr:.4f}, Infty]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q3iJU3cD3DBx"
   },
   "source": [
    "# Marginal Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UJ-p62CD3DBx"
   },
   "outputs": [],
   "source": [
    "important_feats = Z.columns[np.argsort(np.mean(shap_values.values**2, axis=0))[::-1]]\n",
    "important_feats[:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sRb7S8G13DBx"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 7))\n",
    "for it, feature in enumerate(important_feats[:4]):\n",
    "    plt.subplot(2, 2, it + 1)\n",
    "    grid = np.unique(np.percentile(Z[feature], np.arange(0, 105, 5)))\n",
    "    Zpd = pd.DataFrame(np.tile(np.median(Z, axis=0, keepdims=True), (len(grid), 1)),\n",
    "                       columns=Z.columns)\n",
    "    Zpd[feature] = grid\n",
    "\n",
    "    preds = overall_best.predict(Zpd)\n",
    "    preds = preds.flatten()\n",
    "    plt.errorbar(Zpd[feature], preds)\n",
    "    plt.xlabel(feature)\n",
    "    plt.ylabel('Predicted CATE')\n",
    "plt.tight_layout()\n",
    "plt.savefig(f'{data}-marginal-plots.png', dpi=600)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uwU9x-a-3DBx"
   },
   "source": [
    "# Policy Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GiF8DSdl3DBx"
   },
   "source": [
    "Suppose our goal is to estimate the best treatment policy $\\pi: Z \\to \\{0, 1\\}$. The policy gains over no treatment for any policy $\\pi$ can be identified as:\n",
    "\\begin{align}\n",
    "V(\\pi) := E[\\pi(Z)\\, (Y(1) - Y(0))] = E\\left[\\pi(Z)\\, Y^{DR}(g,p)\\right]\n",
    "\\end{align}\n",
    "\n",
    "Moreover, note that optimizing the linear objective $E[\\pi(Z)\\, Y^{DR}(g,p)]$ can also be written as a sample-weighted classification problem, where the goal of $\\pi$ is to much the sign of $Y^{DR}(g,p)$, with sample weights $|Y^{DR}(g,p)|$.\n",
    "\\begin{align}\n",
    "\\text{argmin}_{\\pi} V(\\pi) =~& \\text{argmin}_{\\pi} E\\left[(2\\pi(Z) - 1)\\, Y^{DR}(g,p)\\right]\\\\\n",
    "=~& \\text{argmin}_{\\pi} E\\left[\\left(2\\pi(Z) - 1\\right)\\, \\text{sign}\\left(Y^{DR}(g,p)\\right)\\, \\left|Y^{DR}(g,p)\\right|\\right] \\\\\n",
    "=~& \\text{argmin}_{\\pi} E\\left[1\\left\\{2\\pi(Z) - 1 = \\text{sign}\\left(Y^{DR}(g,p)\\right)\\right\\} \\left|Y^{DR}(g,p)\\right|\\right]\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7XknRKw83DBx"
   },
   "outputs": [],
   "source": [
    "policy = DecisionTreeClassifier(max_depth=2, min_impurity_decrease=1e-3,\n",
    "                                min_samples_leaf=100,\n",
    "                                random_state=random_seed)\n",
    "policy.fit(Z, np.sign(dr_preds - treatment_cost), sample_weight=np.abs(dr_preds - treatment_cost))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BhMAqup63DBx"
   },
   "source": [
    "We can also plot the tree to visualize the partitions and the treatment recommendations. The details that are displayed on each node are also useful in understanding the group average treatment effect for each node. In particular, the information `samples=N`, gives us the size of each node $N$, and the information `value=[A, B]`, then `A` is the sum of the $|Y^{DR}(g,p)|$ for the samples where $Y^{DR}(g,p)<0$ and similarly, `B` is the sum of $|Y^{DR}(g,p)|$ for the samples where $Y^{DR}(g,p)>0$. Thus to get the GATE for each node, we simply do `(B-A)/N`, which would correspond to $\\frac{1}{N}\\sum_{i\\in \\text{node}} Y^{DR}(g,p)$, which is the doubly robust estimate of the GATE for the node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I1-WKGRq3DBx"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(25, 7))\n",
    "plot_tree(policy, filled=True, feature_names=list(Z.columns), impurity=False, label='root',\n",
    "          class_names=['Negative', 'Positive'], fontsize=20)\n",
    "plt.savefig(f'{data}-policy-tree.pdf', dpi=1200)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1A-QzT7y3DBx"
   },
   "source": [
    "#### Evaluating the learned policy out of sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_RchkfkV3DBx"
   },
   "outputs": [],
   "source": [
    "pi = (dr_val - treatment_cost) * policy.predict(Zval)\n",
    "point = np.mean(pi)\n",
    "stderr = np.sqrt(np.var(pi) / pi.shape[0])\n",
    "print(f\"{point:.5f}, {stderr:.5f}, {point - 1.96 * stderr:.5f}, {point + 1.96 * stderr:.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gmTI1fZY3DBz"
   },
   "source": [
    "#### As compared to treating everyone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vmGIusAY3DBz"
   },
   "outputs": [],
   "source": [
    "pi = (dr_val - treatment_cost)\n",
    "point = np.mean(pi)\n",
    "stderr = np.sqrt(np.var(pi) / pi.shape[0])\n",
    "print(f\"{point:.5f}, {stderr:.5f}, {point - 1.96 * stderr:.5f}, {point + 1.96 * stderr:.5f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iMzC_PIN3DBz"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
